#include <cmath>
#include <util/data_type.h>
#include <planner/distance_map_planner/dmp_graph_search.h>

namespace sz {


DMPGraphSearch::DMPGraphSearch(const uint8_t *cMap, short xDim, short yDim, float eps, float cweight, bool verbose) :
    verbose_(verbose),
    xDim_(xDim),
    yDim_(yDim),
    cMap_(cMap),
    eps_(eps),
    dis_map_weight_(cweight) {
  hm_.resize(xDim_ * yDim_);
  seen_.resize(xDim_ * yDim_, false);

  for (short x = -1; x <= 1; x++) {
    for (short y = -1; y <= 1; y++) {
      if (x == 0 && y == 0)
        continue;
      ns_.emplace_back(std::vector<short>{x, y});
    }
  }
}

inline Index DMPGraphSearch::CoordToId(short x, short y) const { return x  + y * xDim_; }

inline bool DMPGraphSearch::IsFree(short x, short y) const {
  return x >= 0 && x < xDim_ && y >= 0 && y < yDim_ && cMap_[CoordToId(x, y)] < val_occ_;
}

inline float DMPGraphSearch::GetHeur(short x, short y) const {
  return eps_ * std::sqrt(float(x - xGoal_) * float(x - xGoal_) + float(y - yGoal_) * float(y - yGoal_));
}

float DMPGraphSearch::Plan(short xStart, short yStart, short xGoal, short yGoal,
                           const std::vector<bool> &in_region) {
  pq_.clear();
  path_.clear();
  hm_.resize(xDim_ * yDim_);
  seen_.resize(xDim_ * yDim_, false);
  in_region_ = in_region;
  if (in_region.empty()) {
    global_ = true;
    if (verbose_)
      printf("global planning!\n");
  } else {
    global_ = false;
    if (verbose_)
      printf("local planning!\n");
  }

  // Set goal
  Index goal_id = CoordToId(xGoal, yGoal);
  xGoal_ = xGoal;
  yGoal_ = yGoal;

  // Set start node
  Index start_id = CoordToId(xStart, yStart);
  DMPStatePtr currNode_ptr = std::make_shared<DMPState>(DMPState(start_id, xStart, yStart));
  currNode_ptr->g = cMap_[start_id];
  currNode_ptr->h = GetHeur(xStart, yStart);

  return Plan(currNode_ptr, start_id, goal_id);
}

float DMPGraphSearch::Plan(DMPStatePtr &currNode_ptr, Index start_id, Index goal_id) {
  // Insert start node
  currNode_ptr->heap_key = pq_.push(currNode_ptr);
  currNode_ptr->opened = true;
  hm_[currNode_ptr->id] = currNode_ptr;
  seen_[currNode_ptr->id] = true;

  int expand_iteration = 0;
  while (true) {
    expand_iteration++;
    // get element with smallest cost
    currNode_ptr = pq_.top();
    pq_.pop();
    currNode_ptr->closed = true; // Add to closed list

    if (currNode_ptr->id == goal_id) {
      if (verbose_)
        printf("Goal Reached!!!!!!\n\n");
      break;
    }

    // printf("expand: %d, %d\n", currNode_ptr->x, currNode_ptr->y);
    std::vector<Index> succ_ids;
    std::vector<float> succ_costs;
    // Get successors
    GetSucc(currNode_ptr, succ_ids, succ_costs);

    // Process successors
    for (int s = 0; s < (int) succ_ids.size(); s++) {
      // see if we can improve the value of succ_state
      DMPStatePtr &child_ptr = hm_[succ_ids[s]];
      float tentative_g_val = currNode_ptr->g + succ_costs[s];

      if (tentative_g_val < child_ptr->g) {
        child_ptr->parent_id = currNode_ptr->id; // Assign new parent
        child_ptr->g = tentative_g_val;          // Update g_val


        if (child_ptr->opened && !child_ptr->closed)        // if currently in OPEN, update
          pq_.increase(child_ptr->heap_key); // update heap

        else if (child_ptr->opened && child_ptr->closed) {          // if currently in CLOSED
          printf("DMP ERROR!\n");
        } else // new node, add to heap
        {
          // printf("add to open set: %d, %d\n", child_ptr->x, child_ptr->y);
          child_ptr->heap_key = pq_.push(child_ptr);
          child_ptr->opened = true;
        }
      } //
    }   // Process successors

    if (pq_.empty()) {
      if (verbose_)
        printf("Priority queue is empty!!!!!!\n\n");
      return std::numeric_limits<float>::infinity();
    }
  }

  if (verbose_) {
    printf("goal g: %f, h: %f!\n", currNode_ptr->g, currNode_ptr->h);
    printf("Expand [%d] nodes!\n", expand_iteration);
  }

  RecoverPath(currNode_ptr, start_id);
  return currNode_ptr->g;
}

void DMPGraphSearch::RecoverPath(DMPStatePtr node, Index start_id) {
  VectorPoint2i path;
  path.emplace_back(node->x, node->y);
  while (node && node->id != start_id) {
    node = hm_[node->parent_id];
    // printf("waypoint g: %f, h: %f!\n", node->g, node->h);
    path.emplace_back(node->x, node->y);
  }
  std::reverse(std::begin(path), std::end(path));
  path.swap(path_);
}

void DMPGraphSearch::GetSucc(const DMPStatePtr &curr, std::vector<Index> &succ_ids,
                             std::vector<float> &succ_costs) {

  for (const auto &d : ns_) {
    short new_x{short(curr->x + d[0])};
    short new_y{short(curr->y + d[1])};
    if (!IsFree(new_x, new_y))
      continue;

    Index new_id = CoordToId(new_x, new_y);
    if (!global_ && !in_region_[new_id])
      continue;

    if (!seen_[new_id]) {
      seen_[new_id] = true;
      hm_[new_id] = std::make_shared<DMPState>(new_id, new_x, new_y);
      hm_[new_id]->h = GetHeur(new_x, new_y);
    }
    succ_ids.push_back(new_id);
    succ_costs.push_back(std::sqrt(float(d[0] * d[0] + d[1] * d[1]) + dis_map_weight_ * (float) (cMap_[new_id])));
  }

}

VectorPoint2i DMPGraphSearch::GetPath() const { return path_; }

std::vector<DMPStatePtr> DMPGraphSearch::GetOpenSet() const {
  std::vector<DMPStatePtr> ss;
  for (const auto &it : hm_) {
    if (it && it->opened && !it->closed)
      ss.push_back(it);
  }
  return ss;
}

std::vector<DMPStatePtr> DMPGraphSearch::GetCloseSet() const {
  std::vector<DMPStatePtr> ss;
  for (const auto &it : hm_) {
    if (it && it->closed)
      ss.push_back(it);
  }
  return ss;
}

std::vector<DMPStatePtr> DMPGraphSearch::GetAllSet() const {
  std::vector<DMPStatePtr> ss;
  for (const auto &it : hm_) {
    if (it)
      ss.push_back(it);
  }
  return ss;
}
}
